package com.tool.engine.export.processor;

import com.tool.engine.export.capacity.CapacityLimit;
import com.tool.engine.export.exception.ExportTaskException;
import com.tool.engine.export.mertics.Collector;
import com.tool.engine.export.output.DataOutputSource;
import com.tool.engine.export.policy.AdjusterPolicy;
import com.tool.engine.export.task.Task;
import com.tool.engine.export.write.WriteTask;
import com.tool.engine.export.write.allocator.DataWriteAllocator;
import lombok.Setter;
import lombok.extern.slf4j.Slf4j;
import org.openjdk.jol.info.GraphLayout;
import org.springframework.util.CollectionUtils;

import java.util.ArrayList;
import java.util.List;
import java.util.concurrent.atomic.AtomicLong;

@Slf4j
public abstract class AbstractProcessor implements ExportTaskProcessor{

    private final ThreadLocal<AtomicLong> dataSizeLocal = new ThreadLocal<>();

    private final ThreadLocal<List<WriteTask>> writeTasksLocal = new ThreadLocal<>();

    @Setter
    private CapacityLimit exportTaskCapacityLimit;

    @Setter
    private DataWriteAllocator dataWriteAllocator;

    @Setter
    private Collector collector;

    @Setter
    private AdjusterPolicy adjusterPolicy;

    public final void doProcessTask(Task<?> task, DataOutputSource dataOutPutSource) throws ExportTaskException {
        List<?> data = null;
        try {
            initData();
            dataOutPutSource.initialize();
            long offset = 0;
            do {
                int batch = adjusterPolicy.getBathSize();
                exportTaskCapacityLimit.tryLock(collector.getAvgMemorySize());
                data = this.fetchData(task, offset, batch);
                if (!CollectionUtils.isEmpty(data)) {
                    long bytesSize = calculateDataSize(data);
                    offset += batch;
                    WriteTask writeTask = createWriteTask(data,dataOutPutSource);
                    dispatcherTask(task.getFileName(), writeTask, bytesSize);
                }
                if (data.size() < batch) {
                    break;
                }
            } while (task.isAll() != null && !task.isAll());
        } catch (Exception e) {
            log.warn("threadName:{} processTask occur exception:{}", Thread.currentThread().getName(), e.getMessage());
            throw new ExportTaskException(e.getMessage());
        } finally {
            data = null; //help gc
        }
    }

    private void dispatcherTask(String fileName, WriteTask writeTask, long bytesSize) {
        dataWriteAllocator.dispatcher(fileName, writeTask, param ->{
            exportTaskCapacityLimit.removeCapacity(bytesSize);
            collector.decrCapacity(bytesSize, exportTaskCapacityLimit.getThreshold());
            return param;
        });
    }

    private WriteTask createWriteTask(List<?> data, DataOutputSource dataOutPutSource) {
        WriteTask writeTask = dataOutPutSource.createWriteTask(data);
        writeTasksLocal.get().add(writeTask);
        return writeTask;
    }

    private long calculateDataSize(List<?> data) {
        long bytesSize = GraphLayout.parseInstance(data).totalSize();
        exportTaskCapacityLimit.addCapacity(bytesSize);
        collector.incrCapacity(data.size(), bytesSize, exportTaskCapacityLimit.getThreshold());
        return bytesSize;
    }

    private List<?> fetchData(Task<?> task, long offset,int batch){
        long fetchDataStartTime = System.currentTimeMillis();
        List<?> data = task.process(batch, offset);
        long fetchDataEndTime = System.currentTimeMillis();
        log.debug("taskName:{} fetchData time:{}ms batch:{}", task.getFileName(), fetchDataEndTime - fetchDataStartTime, batch);
        collector.incrFetchTime(fetchDataEndTime-fetchDataStartTime, data.size());
        return data;
    }

    public final void clear() {
        dataSizeLocal.remove();
        writeTasksLocal.remove();
    }

    private void initData() {
        dataSizeLocal.set(new AtomicLong(0));
        List<WriteTask> writeTasks = new ArrayList<>();
        writeTasksLocal.set(writeTasks);

    }

    @Override
    public boolean isCompleted() {
        List<WriteTask> writeTasks = writeTasksLocal.get();
        for (WriteTask writeTask : writeTasks) {
            if (!writeTask.isSuccess()) {
                log.warn("taskName:{} write failed", writeTask.getTaskName());
                return false;
            }
        }
        return true;
    }
}
